import asyncio
from contextlib import asynccontextmanager
from typing import List
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
import uvicorn
from starlette.websockets import WebSocket, WebSocketDisconnect
from websockets import ConnectionClosedOK
from src.huey_init._huey import huey # 不要删
from src.huey_init import task_status  # 不要删
from src.huey_init.storages import Storage # 不要删
from src.agents.routers import router as agent_router, websocket_endpoint
from src.translation_file.routers import router as tran_router
from src.build_aitemplate.routers import router as aitemplate_router
from src.gen_file.routers import router as genfile_router
from src.edit_file.routers import router as edit_router
from src.find_kv_and_check_word.routers import router as check_router
from src.correct.routers import router as correct_router
from src.tasks.routers import router as task_router
from src.gadget.routers import router as gadget_router
# 维护 WebSocket 连接的列表
websocket_connections: List[WebSocket] = []
from config import REDIS_IP, REDIS_PORT
import redis.asyncio as redis

aio_r = redis.Redis(host=REDIS_IP, port=REDIS_PORT, db=3, health_check_interval=10,
            socket_timeout=10, socket_keepalive=True,
            socket_connect_timeout=10, retry_on_timeout=True)

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load the ML model
    asyncio.create_task(_listener())
    yield
    # Clean up the ML models and release the resources
    websocket_connections = []

app = FastAPI(lifespan=lifespan)
app.include_router(agent_router)
app.include_router(tran_router)
app.include_router(aitemplate_router)
app.include_router(genfile_router)
app.include_router(edit_router)
app.include_router(check_router)
app.include_router(correct_router)
app.include_router(task_router)
app.include_router(gadget_router)
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*']
)


# RabbitMQ 连接参数
# @app.on_event("startup")
# async def startup_event():
#     # 在启动时启动RabbitMQ监听器
#     asyncio.create_task(rabbitmq_listener())

async def _listener():
    """RabbitMQ 消息监听器，将消息发送到 WebSocket 连接"""
    pubsub = aio_r.pubsub()
    channel_name = "channel:1"
    await pubsub.subscribe(channel_name)
    while True:
        message = await pubsub.get_message(ignore_subscribe_messages=True)
        if message is not None:
            data = message['data'].decode("utf-8")
            await broadcast_message(data)


async def broadcast_message(message: str):
    """将消息广播到所有 WebSocket 连接"""
    for websocket in websocket_connections:
        try:
            await websocket.send_text(message)
        except Exception as e:
            print(f"Error sending message to WebSocket: {e}")

@app.get("/")
async def root():
    return {"message": "新都科伦AI文档项目后端"}

@app.websocket("/ws/text_stream")
async def get_gen_string(websocket: WebSocket):
    '''流式输出'''
    await websocket.accept()
    websocket_connections.append(websocket)
    try:
        while True:
            # 等待客户端消息（可以选择忽略接收到的消息）
            await websocket.receive_text()
    except ConnectionClosedOK:
        print("connection closed")
    except WebSocketDisconnect:
        print("disconnect")
    except Exception as e:
        print(f"WebSocket error: {e}")
    finally:
        websocket_connections.remove(websocket)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
    #await sub_gen(websocket, task_id)